package com.chengyu.stock.alg;

import com.google.common.collect.Range;
import org.apache.shardingsphere.api.sharding.standard.PreciseShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.standard.PreciseShardingValue;
import org.apache.shardingsphere.api.sharding.standard.RangeShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.standard.RangeShardingValue;
import org.joda.time.DateTime;

import java.util.Collection;
import java.util.Date;
import java.util.stream.Collectors;

//定义公共的分表算法类：包含精准和范围查询数据
public class CommonAlg4Tb implements PreciseShardingAlgorithm<Date>, RangeShardingAlgorithm<Date> {
    //精准查询，eg:select * from stock_rt_info where cur_time='2022-12-03 09:40:00'  (in =)
    @Override
    public String doSharding(Collection<String> tbNames, PreciseShardingValue<Date> shardingValue) {
        //获取逻辑表
        String logicTableName = shardingValue.getLogicTableName();
        //获取分片字段名称
        String columnName = shardingValue.getColumnName();
        //获取精准查询的条件值
        Date curTime = shardingValue.getValue();
        //根据查询的日期时间中的年份获取对应的数据源
        String curYear = new DateTime(curTime).toString("yyyyMM");
        String name = tbNames.stream().filter(tbName -> tbName.endsWith(curYear)).findFirst().get();
        return name;
    }

    @Override
    public Collection<String> doSharding(Collection<String> tbNames, RangeShardingValue<Date> shardingValue) {
        //获取范围封装对象
        Range<Date> valueRange = shardingValue.getValueRange();
        //判断上限和下限
        //过滤出小于等于起始年份的数据源
        if (valueRange.hasUpperBound()) {
            //获取上限值
            Integer yearMonth = Integer.valueOf(new DateTime(valueRange.upperEndpoint()).toString("yyyyMM"));
            //过滤数据源中后缀小于等于上限值的数据源集合
            tbNames = tbNames.stream().filter(tbName -> Integer.valueOf(tbName.substring(tbName.lastIndexOf("_") + 1)) <= yearMonth).collect(Collectors.toList());
        }
        //过滤出大于等于起始年份的数据源
        if (valueRange.hasLowerBound()) {
            //获取上限值
            Integer yearMonth = Integer.valueOf(new DateTime(valueRange.lowerEndpoint()).toString("yyyyMM"));//2022
            //只需过滤出数据源中后缀小于等于上限值的数据源集合集合
            tbNames = tbNames.stream().filter(tbName -> Integer.valueOf(tbName.substring(tbName.lastIndexOf("_") + 1)) >= yearMonth).collect(Collectors.toList());
        }
        return tbNames;
    }
}
